

from torchvision import datasets, transforms
from torch.utils import data
from torch.nn import functional as F

from utils import q_sample as forward_process
from model import Unet


transform = transforms.Compose([
    transforms.Resize(64, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.RandomHorizontalFlip(0.4),
    transforms.ToTensor(),  # scale to [0,1]
    transforms.Lambda(lambda t: (t * 2) - 1)  # scale to [-1,1]
])


dataset = datasets.MNIST("X:\Machine_Learning\datasets\MNIST_\\",
                         download=False,
                         transform=transform)

dataloader = data.DataLoader(dataset, batch_size=32,pin_memory=True)
model = Unet(isGrayImage=True)
# To implenmente such a model , we need a Scheduler to add noise to images.
# And a neuron network to predict the noise.

certerion = F.mse_loss

for x,_ in dataloader:
    
    break;

